"""
Text-to-Speech (TTS) 服务模块
使用 Coqui TTS 进行文字转语音
"""

import os
import logging
import hashlib
import tempfile
from typing import Tuple, Optional

# 设置TTS缓存目录到用户目录，避免权限问题
os.environ['XDG_CACHE_HOME'] = os.path.expanduser('~/.cache')
os.environ['XDG_DATA_HOME'] = os.path.expanduser('~/.local/share')
os.environ['TTS_CACHE_DIR'] = os.path.expanduser('~/.cache/tts')
os.environ['TORCH_HOME'] = os.path.expanduser('~/.cache/torch')
os.environ['HF_HOME'] = os.path.expanduser('~/.cache/huggingface')

try:
    from TTS.api import TTS
    TTS_AVAILABLE = True
except ImportError:
    TTS_AVAILABLE = False
    logging.warning("Coqui TTS 未安装，文字转语音功能不可用。请运行: pip install TTS")

from . import config
from .cache_service import cache_service, CacheKeys


class TTSService:
    """Coqui TTS 服务类"""
    
    def __init__(self):
        self.tts_models = {}  # 缓存已加载的模型
        self.current_model = None
        self.current_model_type = None
        
        # 定义可用的模型
        self.available_models = {
            "bark": ("tts_models/multilingual/multi-dataset/bark", "Bark - 最自然的多语言TTS"),
            "your_tts": ("tts_models/multilingual/multi-dataset/your_tts", "YourTTS - 高质量多语言模型"),
            "xtts_v2": ("tts_models/multilingual/multi-dataset/xtts_v2", "XTTS v2 - 最先进的克隆语音"),
            "chinese": ("tts_models/zh-CN/baker/tacotron2-DDC-GST", "Baker - 中文专用模型"),
            "english_vits": ("tts_models/en/ljspeech/vits", "VITS - 高质量英文模型"),
            "english": ("tts_models/en/ljspeech/tacotron2-DDC", "Tacotron2 - 稳定英文模型")
        }
    
    def load_model(self, model_type: str):
        """加载指定的TTS模型"""
        if not TTS_AVAILABLE:
            logging.error("TTS库未安装，无法加载TTS模型")
            return False
        
        if model_type not in self.available_models:
            logging.error(f"不支持的模型类型: {model_type}")
            return False
        
        # 如果模型已经加载且是当前模型，直接返回
        if model_type == self.current_model_type and self.current_model is not None:
            return True
        
        # 检查是否已缓存
        if model_type in self.tts_models:
            self.current_model = self.tts_models[model_type]
            self.current_model_type = model_type
            logging.info(f"切换到已缓存的TTS模型: {model_type}")
            return True
        
        # 加载新模型
        model_name, description = self.available_models[model_type]
        try:
            logging.info(f"正在加载 {description}: {model_name}")
            tts_model = TTS(model_name)
            self.tts_models[model_type] = tts_model
            self.current_model = tts_model
            self.current_model_type = model_type
            logging.info(f"TTS模型加载成功: {description}")
            return True
        except Exception as e:
            logging.error(f"{description} 加载失败: {e}")
            return False
    
    def _load_fallback_model(self):
        """加载备用模型"""
        # 尝试顺序：中文模型 -> 英文模型 -> 其他
        fallback_order = ["chinese", "english", "english_vits", "your_tts", "bark"]
        
        for model_type in fallback_order:
            if self.load_model(model_type):
                logging.info(f"使用备用模型: {model_type}")
                return True
        
        logging.error("所有备用模型都加载失败")
        return False
    
    def _apply_male_voice_effect(self, audio_path: str):
        """对音频应用男声效果（降低音调）"""
        try:
            import librosa
            import soundfile as sf
            
            # 读取音频
            audio, sr = librosa.load(audio_path, sr=None)
            
            # 降低音调来模拟男声（-3到-5个半音）
            # pitch_shift参数：负值降低音调，正值提高音调
            male_audio = librosa.effects.pitch_shift(audio, sr=sr, n_steps=-4)
            
            # 保存处理后的音频
            sf.write(audio_path, male_audio, sr)
            logging.info("男声音调处理完成")
            
        except ImportError:
            logging.warning("librosa或soundfile库未安装，无法进行男声后处理。可以运行: pip install librosa soundfile")
        except Exception as e:
            logging.warning(f"男声音调处理失败: {e}")
    
    def _apply_voice_speed_adjustment(self, audio_path: str, speed_factor: float = 1.0):
        """调整音频播放速度"""
        try:
            import librosa
            import soundfile as sf
            
            # 读取音频
            audio, sr = librosa.load(audio_path, sr=None)
            
            # 调整速度
            adjusted_audio = librosa.effects.time_stretch(audio, rate=speed_factor)
            
            # 保存处理后的音频
            sf.write(audio_path, adjusted_audio, sr)
            logging.info(f"音频速度调整完成: {speed_factor}x")
            
        except ImportError:
            logging.warning("librosa或soundfile库未安装，无法进行速度调整")
        except Exception as e:
            logging.warning(f"音频速度调整失败: {e}")
    
    def is_available(self) -> bool:
        """检查TTS服务是否可用"""
        return TTS_AVAILABLE and self.current_model is not None
    
    def generate_speech(self, text: str, language: str = "zh", speaker: str = "female", model_type: str = "bark") -> Optional[str]:
        """
        将文字转换为语音
        
        Args:
            text: 要转换的文字内容
            language: 语言代码 (zh, en, ja, ko, es, fr, de)
            speaker: 说话人类型 (female, male)
            model_type: TTS模型类型
            
        Returns:
            生成的音频文件路径，失败返回None
        """
        # 确保指定的模型已加载
        if not self.load_model(model_type):
            logging.warning(f"无法加载指定模型 {model_type}，尝试使用备用模型")
            if not self._load_fallback_model():
                logging.error("TTS服务不可用")
                return None
        
        if not text.strip():
            logging.warning("文字内容为空，跳过TTS生成")
            return None
        
        # 生成缓存键（包含模型类型）
        text_hash = hashlib.md5(f"{text}_{language}_{speaker}_{model_type}".encode()).hexdigest()
        cache_key = f"{CacheKeys.TTS_AUDIO}:{text_hash}"
        
        # 尝试从缓存获取
        cached_audio_path = cache_service.get(cache_key)
        if cached_audio_path and os.path.exists(cached_audio_path):
            logging.info("从缓存获取TTS音频文件")
            return cached_audio_path
        
        try:
            logging.info(f"正在生成TTS语音: 模型={model_type}, 语言={language}, 说话人={speaker}, 文字长度={len(text)}")
            
            # 映射语言代码到TTS支持的语言
            language_map = {
                "zh": "zh-cn",
                "en": "en", 
                "ja": "ja",
                "ko": "ko",
                "es": "es",
                "fr": "fr",
                "de": "de"
            }
            
            tts_language = language_map.get(language, "en")
            
            # 创建临时音频文件
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False, dir=config.WORK_DIR) as temp_file:
                temp_audio_path = temp_file.name
            
            # 根据模型类型使用不同的参数
            if model_type == "bark":
                # Bark模型：支持多语言，自然语音
                # 根据语言和speaker选择voice preset
                if language == "zh":
                    voice_preset = "v2/zh_speaker_7" if speaker == "female" else "v2/zh_speaker_8"
                else:
                    voice_preset = "v2/en_speaker_6" if speaker == "female" else "v2/en_speaker_9"
                
                self.current_model.tts_to_file(
                    text=text,
                    file_path=temp_audio_path,
                    voice_preset=voice_preset
                )
            elif model_type == "your_tts":
                # YourTTS：多语言模型
                self.current_model.tts_to_file(
                    text=text,
                    file_path=temp_audio_path,
                    language=tts_language
                )
            elif model_type == "xtts_v2":
                # XTTS v2：最先进但需要speaker
                try:
                    self.current_model.tts_to_file(
                        text=text,
                        file_path=temp_audio_path,
                        language=tts_language,
                        split_sentences=True
                    )
                except Exception as e:
                    # 如果需要speaker，使用内置speaker
                    logging.warning(f"XTTS v2 需要speaker参数: {e}")
                    if hasattr(self.current_model, 'speakers') and self.current_model.speakers:
                        speaker_name = self.current_model.speakers[0]
                        self.current_model.tts_to_file(
                            text=text,
                            file_path=temp_audio_path,
                            speaker=speaker_name,
                            language=tts_language
                        )
                    else:
                        raise e
            elif model_type == "chinese":
                # 中文专用模型（Baker模型是单speaker，需要后处理模拟男女声）
                self.current_model.tts_to_file(
                    text=text,
                    file_path=temp_audio_path
                )
                
                # 如果选择了男声，对音频进行后处理降低音调
                if speaker == "male":
                    logging.info("对中文模型音频进行男声后处理...")
                    self._apply_male_voice_effect(temp_audio_path)
            elif model_type in ["english", "english_vits"]:
                # 英文模型
                self.current_model.tts_to_file(
                    text=text,
                    file_path=temp_audio_path
                )
            else:
                # 默认尝试
                self.current_model.tts_to_file(
                    text=text,
                    file_path=temp_audio_path
                )
            
            # 验证生成的音频文件
            if not os.path.exists(temp_audio_path) or os.path.getsize(temp_audio_path) == 0:
                logging.error("TTS生成的音频文件无效")
                return None
            
            # 移动到最终位置
            final_audio_path = os.path.join(config.WORK_DIR, f"tts_audio_{text_hash}.wav")
            os.rename(temp_audio_path, final_audio_path)
            
            # 缓存结果
            cache_service.set(cache_key, final_audio_path, ttl=7200)  # 缓存2小时
            
            logging.info(f"TTS语音生成完成: {final_audio_path}")
            return final_audio_path
            
        except Exception as e:
            logging.error(f"TTS语音生成失败: {e}")
            # 清理临时文件
            if 'temp_audio_path' in locals() and os.path.exists(temp_audio_path):
                os.remove(temp_audio_path)
            return None
    
    def create_audio_from_text(self, text: str, language: str = "zh", speaker: str = "female", model_type: str = "bark") -> Optional[str]:
        """
        为长文本创建音频，自动分段处理
        
        Args:
            text: 要转换的文字内容
            language: 语言代码
            speaker: 说话人类型
            model_type: TTS模型类型
            
        Returns:
            最终合成的音频文件路径
        """
        if not text.strip():
            return None
        
        # 分段处理长文本（每段最多500字符）
        max_segment_length = 500
        segments = []
        
        # 按句子分割
        import re
        sentences = re.split(r'[。！？.!?]', text)
        
        current_segment = ""
        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue
                
            if len(current_segment + sentence) <= max_segment_length:
                current_segment += sentence + "。"
            else:
                if current_segment:
                    segments.append(current_segment.strip())
                current_segment = sentence + "。"
        
        if current_segment:
            segments.append(current_segment.strip())
        
        if not segments:
            return None
        
        logging.info(f"文本分为 {len(segments)} 段进行TTS处理")
        
        # 为每段生成音频
        audio_files = []
        for i, segment in enumerate(segments):
            logging.info(f"处理第 {i+1}/{len(segments)} 段...")
            audio_path = self.generate_speech(segment, language, speaker, model_type)
            if audio_path:
                audio_files.append(audio_path)
            else:
                logging.warning(f"第 {i+1} 段TTS生成失败，跳过")
        
        if not audio_files:
            logging.error("所有段落的TTS生成都失败")
            return None
        
        # 如果只有一个音频文件，直接返回
        if len(audio_files) == 1:
            return audio_files[0]
        
        # 合并多个音频文件
        return self._merge_audio_files(audio_files, text)
    
    def _merge_audio_files(self, audio_files: list, original_text: str) -> Optional[str]:
        """合并多个音频文件"""
        try:
            import moviepy.editor as mp
            
            logging.info(f"正在合并 {len(audio_files)} 个音频文件...")
            
            # 加载所有音频片段
            audio_clips = []
            for audio_file in audio_files:
                if os.path.exists(audio_file):
                    clip = mp.AudioFileClip(audio_file)
                    audio_clips.append(clip)
                    # 在片段之间添加短暂停顿（0.5秒）
                    silence = mp.AudioClip(lambda t: 0, duration=0.5).set_fps(22050)
                    audio_clips.append(silence)
            
            if not audio_clips:
                return None
            
            # 移除最后一个停顿
            if len(audio_clips) > 1:
                audio_clips.pop()
            
            # 合并音频
            final_audio = mp.concatenate_audioclips(audio_clips)
            
            # 生成输出文件名
            text_hash = hashlib.md5(original_text.encode()).hexdigest()[:8]
            output_path = os.path.join(config.WORK_DIR, f"tts_merged_{text_hash}.wav")
            
            # 写入文件
            final_audio.write_audiofile(output_path, verbose=False, logger=None)
            
            # 清理临时资源
            for clip in audio_clips:
                if hasattr(clip, 'close'):
                    clip.close()
            final_audio.close()
            
            # 删除临时分段文件
            for audio_file in audio_files:
                try:
                    os.remove(audio_file)
                except:
                    pass
            
            logging.info(f"音频合并完成: {output_path}")
            return output_path
            
        except Exception as e:
            logging.error(f"音频合并失败: {e}")
            return None


# 全局TTS服务实例
tts_service = TTSService()


def create_audio_from_text(text: str, language: str = "zh", speaker: str = "female", model_type: str = "bark") -> Optional[str]:
    """
    便捷函数：将文字转换为音频
    
    Args:
        text: 要转换的文字内容
        language: 语言代码 (zh, en, ja, ko, es, fr, de)  
        speaker: 说话人类型 (female, male)
        model_type: TTS模型类型
        
    Returns:
        生成的音频文件路径，失败返回None
    """
    return tts_service.create_audio_from_text(text, language, speaker, model_type)


def is_tts_available() -> bool:
    """检查TTS服务是否可用"""
    return tts_service.is_available()
